from dataclasses import dataclass

@dataclass
class Nougat:
    model_path: str = '/home/work/model/torch_hub/nougat-0.1.0-base'
    batch_size: int = 8
    # start from page 1
    stop_page: int = 16
    save_dir: str = '_data/pdfs'
    gpu_id: str = '5'
    base_url: str = 'http://127.0.0.1:8503/predict/'

@dataclass
class LLM:
    #model_path: str = '/home/work/model/ECNU-SEA/SEA-E'
    model_path: str = '/home/work/model/TEST/TEST'
    model: str = None
    gpu_id: str = '4'
    openai_base_url: str = 'http://127.0.0.1:8765/api/v1/'
    openai_api_key: str = 'EMPTY'

    def __post_init__(self):
        self.model = self.model_path.split('/')[-1]

@dataclass
class Base:
    app_port: int = 8070
    https_proxy: str = '[TO BE ADDED]'
    log_dir: str = '_log'
    data_dir: str = '_data'

